package main

import (
	"bytes"
	"context"
	"errors"
	"fmt"
	"github.com/klauspost/compress/zstd"
	"io"
	"net"
	"os"
	"path/filepath"
	"time"

	"github.com/ipfs/boxo/ipld/merkledag"
	"github.com/ipfs/go-cid"
	mh "github.com/multiformats/go-multihash"

	"github.com/vmihailenco/msgpack/v5"
)

const maxBlockSize = 1024 * 1024 // 1MB

// Define a struct to represent a file block
type FileBlock struct {
	Index int    // Index of the block within the file
	Data  []byte // Data of the block
	Cid   string // CID of the block
}

// Define a function to send a file block
func sendBlock(w io.Writer, block *FileBlock) error {
	//msgEnc := msgpack.NewEncoder(w)
	//err := msgEnc.Encode(block)
	marshal, err := msgpack.Marshal(block)
	if err != nil {
		return err
	}
	fmt.Println("sending block: ", block.Index)
	_, err = w.Write(marshal)
	if err != nil {
		return err
	}
	return err
}

// Define a function to receive a file block
func recvBlock(r io.Reader, expectedIndex int) (*FileBlock, error) {
	fb := new(FileBlock)
	//msgDec := msgpack.NewDecoder(r)
	//err := msgDec.Decode(fb)
	buf := make([]byte, maxBlockSize)
	fmt.Println("recving block: ", expectedIndex)
	n, err := r.Read(buf)
	if err != nil {
		return nil, err
	}
	err = msgpack.Unmarshal(buf[:n], fb)
	if err != nil {
		return nil, err
	}
	if fb.Index != expectedIndex {
		return nil, fmt.Errorf("received block with unexpected index %d (expected %d)", fb.Index, expectedIndex)
	}
	if err != nil && err != io.ErrUnexpectedEOF {
		return nil, fmt.Errorf("read data: %s", err.Error()) //
	}
	// 打印得到的文件块的内容
	//fmt.Println(string(fb.Data))

	datalen := len(fb.Data)
	if datalen == 0 {
		return nil, nil
	}
	return fb, nil
}

// New function to send the root CID and file size
func sendFileSize(w io.Writer, fileSize uint64) error {
	//msgEnc := msgpack.NewEncoder(w)
	//err := msgEnc.Encode(fileSize)
	marshal, err := msgpack.Marshal(fileSize)
	if err != nil {
		return err
	}
	_, err = w.Write(marshal)
	if err != nil {
		return err
	}
	return nil
}

// New function to send the root CID and file size
func sendRootCID(w io.Writer, rootCID cid.Cid) error {
	//msgEnc := msgpack.NewEncoder(w)
	//err := msgEnc.Encode(rootCID.Bytes())
	fmt.Print("sending rootcid  ------->  ")
	marshal, err := msgpack.Marshal(rootCID.Bytes())
	if err != nil {
		return err
	}
	_, err = w.Write(marshal)
	fmt.Println(" sended rootcid")
	if err != nil {
		return err
	}
	return nil
}

// New function to receive the root CID and file size 勾
func recvFileSize(r io.Reader) (uint64, error) {
	//msgDec := msgpack.NewDecoder(r)
	var fileSize uint64
	//err := msgDec.Decode(&fileSize)
	buf := make([]byte, maxBlockSize)
	n, err := r.Read(buf)
	if err != nil {
		return 0, nil
	}
	err = msgpack.Unmarshal(buf[:n], &fileSize)
	if err != nil {
		return 0, err
	}

	if err != nil {
		return 0, fmt.Errorf("decode file size: %s", err.Error())
	}
	return fileSize, nil
}

// New function to receive the root CID and file size 勾
func recvRootCID(r io.Reader) (cid.Cid, error) {
	//msgDec := msgpack.NewDecoder(r)
	var cidBytes []byte
	//err := msgDec.Decode(&cidBytes)
	buf := make([]byte, maxBlockSize)
	fmt.Print("recving rootcid  -------> ")
	n, err := r.Read(buf)
	fmt.Println(" recved rootcid")
	if err != nil {
		return cid.Cid{}, err
	}
	err = msgpack.Unmarshal(buf[:n], &cidBytes)
	if err != nil {
		return cid.Cid{}, err
	}
	if err != nil {
		return cid.Undef, fmt.Errorf("decode CID: %s", err.Error())
	}
	rootCID, err := cid.Cast(cidBytes)
	if err != nil {
		return cid.Undef, fmt.Errorf("invalid CID: %s", err.Error())
	}
	return rootCID, nil
}

type cidBuilder struct {
	Codec    uint64
	MhType   uint64
	MhLength int // MhLength <= 0 means the default length
}

func newCidBuilder(codec uint64, mhType uint64) *cidBuilder {
	return &cidBuilder{Codec: codec, MhType: mhType}
}

func (c cidBuilder) Sum(data []byte) (cid.Cid, error) {
	mhLen := c.MhLength
	if mhLen <= 0 {
		mhLen = -1
	}
	hash, err := mh.Sum(data, c.MhType, mhLen)
	if err != nil {
		return cid.Undef, err
	}
	return cid.NewCidV1(c.Codec, hash), nil
}

func (c cidBuilder) GetCodec() uint64 {
	return c.Codec
}

func (c cidBuilder) WithCodec(u uint64) cid.Builder {
	c.Codec = u
	return c
}

// Modified sendfile function
func sendfile(rwc io.ReadWriteCloser, fn string) error {
	defer func(rwc io.ReadWriteCloser) {
		err := rwc.Close()
		if err != nil {

		}
	}(rwc)
	filename := filepath.Join(sharePath, fn)
	info, err := os.Stat(filename)
	if err != nil {
		return err
	}
	fp, _ := os.Open(filename)
	defer func(fp *os.File) {
		err := fp.Close()
		if err != nil {

		}
	}(fp)
	// TODO 从数据库中得到文件路径从而获取文件句柄，以及对应根cid。目前是直接先得到了句柄，所以这一步暂时可以放一下

	_, err = fp.Seek(0, io.SeekStart)
	if err != nil {
		return err
	}

	blockNum := getFileBlocksNum(uint64(info.Size()), maxBlockSize)
	err = sendFileBlocksNum(rwc, blockNum)
	if err != nil {
		return err
	}
	//err = sendFileSize(rwc, uint64(info.Size()))
	//if err != nil {
	//	return err
	//}

	blockIndex := 0
	///////////////////////////////////////////////////////////////
	root := merkledag.NodeWithData(nil)
	cidBuilder := newCidBuilder(cid.DagProtobuf, mh.BLAKE3)
	err = root.SetCidBuilder(cidBuilder)
	if err != nil {
		return err
	}

	///////////////////////////////////////////////////////////////////
	for {
		data := make([]byte, maxBlockSize)
		n, err := fp.Read(data)

		if err != nil && err != io.EOF {
			return err
		}
		if n == 0 || uint64(blockIndex) == blockNum {

			break
		}
		if n < maxBlockSize {
			data = data[:n]
		}
		// 经过上面的操作得到一个文件块

		block := &FileBlock{Index: blockIndex, Data: data}
		//cid, _ := cidBuilder.Sum(data)

		//////////////////////////////////////
		node, _ := merkledag.NewRawNodeWPrefix(data, cidBuilder)

		block.Cid = node.Cid().String()
		//fmt.Println("put: ", block.Cid)
		if err != nil {
			return err
		}
		err = sendBlock(rwc, block)
		if err != nil {
			return err
		}
		//fmt.Println("node.cid: ", node.Cid().String())
		//fmt.Println("cid     : ", cid.String())
		root.AddNodeLink(node.Cid().String(), node)
		//////////////////////////////////////////////////////////////
		blockIndex++
	}
	// 创建一个带有超时的 context
	ctx, cancel := context.WithTimeout(context.Background(), 1*time.Second)
	defer cancel()
	// 发送根CID并等待回应
	err = sendRootCIDWithRetry(ctx, rwc, root.Cid(), maxBlockSize)
	if err != nil {
		return err
	}
	////////////////////////////
	return nil
}
func sendRootCIDWithRetry(ctx context.Context, rwc io.ReadWriteCloser, rootCid cid.Cid, maxBlockSize int) error {
	buf := make([]byte, maxBlockSize)
	var msg string
	retryCount := 0

	for {
		// 发送根CID
		err := sendRootCID(rwc, rootCid)
		if err != nil {
			return err
		}

		// 等待接收方回应
		fmt.Println("waiting for receiver call done")
		err = rwc.(interface {
			SetReadDeadline(time.Time) error
		}).SetReadDeadline(time.Now().Add(1 * time.Second)) // 设置读取超时时间为3秒
		if err != nil {
			return err
		}
		n, err := readWithTimeout(ctx, rwc, buf) // 使用带有超时的 readWithTimeout 函数读取数据
		if err != nil {
			if netErr, ok := err.(net.Error); ok && netErr.Timeout() {
				// 如果读取超时，则重新发送根CID
				if retryCount < 3 { // 最多重试3次
					fmt.Println("timeout, retrying...")
					retryCount++
					continue
				} else {
					return errors.New("timeout, max retries exceeded")
				}
			}
			return err
		}
		err = msgpack.Unmarshal(buf[:n], &msg)
		if err != nil {
			return err
		}
		if msg == "done" {
			fmt.Println("trans finish")
			return nil
		} else {
			fmt.Println("not finish trans, retrying...")
			if retryCount < 3 { // 最多重试3次
				retryCount++
				continue
			} else {
				return errors.New("not finish trans, max retries exceeded")
			}
		}
	}
}
func readWithTimeout(ctx context.Context, r io.Reader, buf []byte) (int, error) {
	done := make(chan error, 1)

	go func() {
		_, err := r.Read(buf)
		if err != nil {
			done <- err
		} else {
			done <- nil
		}
	}()

	select {
	case <-ctx.Done():
		return 0, ctx.Err()
	case err := <-done:
		if err != nil {
			return 0, err
		}
		return len(buf), nil
	}
}

// Modified recvfile function
func recvfile(rw io.ReadWriter, filename string) error {
	// Receive root CID and file size

	blocksnum, err := recvFileBlocksNum(rw)
	if err != nil {
		return err
	}
	//fileSize, err := recvFileSize(rw)
	//if err != nil {
	//	return err
	//}
	//fmt.Println("fileSize: ", fileSize)

	fmt.Println("recvblocknum: ", blocksnum)
	fp, err := os.Create("../recv/" + filename)
	if err != nil {
		return err
	}
	defer fp.Close()

	//stat, err := fp.Stat()
	//if err != nil {
	//	return err
	//}

	// Receive each block and write it to the file
	blockIndex := 0
	for {
		block, err := recvBlock(rw, blockIndex)
		if err != nil && err != io.EOF {
			return err
		}

		if err != io.EOF {
			_, err1 := fp.Write(block.Data)
			if err1 != nil {
				return err1
			}
		}
		//fmt.Println(stat.Size(), "size")
		if err == io.EOF || uint64(blockIndex)+1 == blocksnum {
			fmt.Println("has get all blocks:", blockIndex+1) // 能走到这里
			break
		}
		blockIndex++
	}
	//fmt.Println("test")
	exrootCID, err := recvRootCID(rw)
	if err != nil {
		return err
	}
	fp.Seek(0, io.SeekStart)

	err, rootCID := GetCid(fp)
	if err != nil {
		return err
	}
	if !exrootCID.Equals(rootCID) {
		fmt.Println("exrootCiD: ", exrootCID.String())
		fmt.Println("rootCid  : ", rootCID.String())
		return errors.New("received file CID does not match the expected CID")
	} else {
		msg := "done"
		marshal, err := msgpack.Marshal(msg)
		if err != nil {
			return err
		}
		_, err = rw.Write(marshal)
		if err != nil {
			return err
		}
	}
	return nil
}

func sendFileBlocksNum(w io.Writer, num uint64) error {
	msgEnc := msgpack.NewEncoder(w)
	err := msgEnc.Encode(num)
	//fmt.Println("blockNum: ", num)
	if err != nil {
		return err
	}
	return nil
}

func recvFileBlocksNum(r io.Reader) (uint64, error) {
	msgDec := msgpack.NewDecoder(r)
	var num uint64
	err := msgDec.Decode(&num)
	if err != nil {
		return 0, fmt.Errorf("decode fileblocksnum: %s", err.Error())
	}
	return num, nil
}

func getFileBlocksNum(fileSize uint64, blockSize uint64) uint64 {
	quotient := fileSize / blockSize
	remainder := fileSize % blockSize
	if remainder != 0 {
		return quotient + 1
	}
	return quotient
}

func GetCid(fp *os.File) (error, cid.Cid) {
	root := merkledag.NodeWithData(nil)
	cidBuilder := newCidBuilder(cid.DagProtobuf, mh.BLAKE3)
	root.SetCidBuilder(cidBuilder)
	for {
		data := make([]byte, maxBlockSize)
		n, err := fp.Read(data)

		if err != nil && err != io.EOF {
			return err, cid.Cid{}
		}
		if n == 0 {
			//params.Dagserv.Add(context.Background(), root)
			//fmt.Println(root.Cid().String())
			break
		}
		if n < maxBlockSize {
			data = data[:n]
		}
		// 经过上面的操作得到一个文件块
		node, _ := merkledag.NewRawNodeWPrefix(data, cidBuilder)
		root.AddNodeLink(node.Cid().String(), node)
	}
	return nil, root.Cid()
}

// 压缩传输
func dataZstd(data []byte) ([]byte, error) {
	buf := bytes.NewBufferString("")
	zw, err := zstd.NewWriter(buf, zstd.WithEncoderLevel(zstd.SpeedBestCompression))
	if err != nil {
		return nil, err
	}
	_, err = zw.Write(data)
	if err != nil {
		return nil, err
	}
	zw.Close()
	return buf.Bytes(), nil
}

func dataUnzstd(data []byte) ([]byte, error) {
	buf := bytes.NewBuffer(data)
	zr, err := zstd.NewReader(buf)
	if err != nil {
		return nil, err
	}
	out := bytes.NewBufferString("")
	_, err = out.ReadFrom(zr)
	if err != nil {
		return nil, fmt.Errorf("Not EOF:%s", err.Error())
	}
	zr.Close()
	return out.Bytes(), nil
}
